package com.ds.infra.mq.core.api;

import com.ds.infra.mq.core.api.exception.ProduceException;
import com.ds.infra.mq.core.protocol.Message;
import com.ds.infra.mq.core.protocol.MessageHeader;
import com.ds.infra.mq.core.protocol.MessageType;
import com.ds.infra.mq.core.protocol.ProtocolCodec;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.util.concurrent.atomic.AtomicLong;

/**
 * Producer的实现类，用于向Broker发送消息
 *
 * @author WeiShaoying
 */
@Slf4j
public class MQProducer implements Producer {

    private final String brokerHost;
    private final int brokerPort;
    private final String clientId;
    private final AtomicLong requestIdGenerator = new AtomicLong(0);

    private EventLoopGroup group;
    private Channel channel;

    public MQProducer(String brokerHost, int brokerPort, String clientId) {
        this.brokerHost = brokerHost;
        this.brokerPort = brokerPort;
        this.clientId = clientId;
    }

    public void connect() throws InterruptedException {
        group = new NioEventLoopGroup(1);
        Bootstrap bootstrap = new Bootstrap();
        bootstrap.group(group)
                .channel(NioSocketChannel.class)
                .handler(new ChannelInitializer<SocketChannel>() {
                    @Override
                    protected void initChannel(SocketChannel ch) {
                        ch.pipeline()
                                .addLast(new ProtocolCodec.Decoder())
                                .addLast(new ProtocolCodec.Encoder());
                    }
                });

        ChannelFuture future = bootstrap.connect(brokerHost, brokerPort).sync();
        channel = future.channel();
        log.info("Producer connected to {}:{}", brokerHost, brokerPort);
    }

    @Override
    public void send(String topic, byte[] message) throws ProduceException {
        send(topic, 0, message);
    }

    @Override
    public void send(String topic, int partition, byte[] message) throws ProduceException {
        MessageHeader header = new MessageHeader();
        header.setMessageType(MessageType.PRODUCE_REQUEST.getCode());
        header.setRequestId(requestIdGenerator.incrementAndGet());
        header.setClientId(clientId);

        // 构建符合格式的消息体：topic:partition|messageBody
        String formattedMessage = topic + ":" + partition + "|" + new String(message, StandardCharsets.UTF_8);

        Message msg = new Message();
        msg.setHeader(header);
        msg.setBody(formattedMessage.getBytes(StandardCharsets.UTF_8));

        try {
            ChannelFuture future = channel.writeAndFlush(msg).sync();
            if (!future.isSuccess()) {
                throw new ProduceException("Failed to send message", future.cause());
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new ProduceException("Send interrupted", e);
        }
    }

    @Override
    public void close() {
        if (channel != null) {
            channel.close();
        }
        if (group != null) {
            group.shutdownGracefully();
        }
        log.info("Producer closed");
    }
}
